import struct
import zlib
import image_restorer


class PngChunk:
    def __init__(self, image_data, offset):
        self.offset = offset
        self.image_data = image_data
        self.init_chunk()

    def init_chunk(self):
        """
        初始化chunk信息
        """
        # 读取chunk长度(4字节)
        self.length = struct.unpack(
            ">I", self.image_data[self.offset : self.offset + 4]
        )[0]
        # 读取chunk类型(4字节)
        self.type = self.image_data[self.offset + 4 : self.offset + 8]
        # 读取chunk数据
        self.data = self.image_data[self.offset + 8 : self.offset + 8 + self.length]
        # 读取CRC(4字节)
        self.crc = struct.unpack(
            ">I",
            self.image_data[
                self.offset + 8 + self.length : self.offset + 12 + self.length
            ],
        )[0]

    def set_data(self, data: bytes):
        """
        设置chunk数据
        """
        self.data = data
        new_data = bytearray(self.image_data)
        # 修改数据段
        new_data[self.offset + 8 : self.offset + 8 + self.length] = data
        self.image_data = bytes(new_data)
        # print(self.image_data)

    def check_crc(self):
        """
        校验CRC
        """
        # 计算CRC
        crc = zlib.crc32(self.type + self.data) & 0xFFFFFFFF
        # 比较CRC
        if crc == self.crc:
            return True
        else:
            return False

    def get_chunk_type(self):
        return self.type.decode()


class BaseChunkRestorer:
    """
    chunk修复
    """

    def __init__(self, chunk: PngChunk):
        self.chunk = chunk

    def repair(self):
        pass

class IhdrChunkRestorer(BaseChunkRestorer):
    def __init__(self, png_chunk: PngChunk):
        if png_chunk.get_chunk_type() == "IHDR":
            self.png_chunk = png_chunk
            self.load_chunk_data()
        else:
            print("Not IHDR chunk")

    def load_chunk_data(self):
        """
        加载IHDR chunk数据
        """
        self.width = struct.unpack(">I", self.png_chunk.data[:4])[0]
        self.height = struct.unpack(">I", self.png_chunk.data[4:8])[0]
        self.bit_depth = self.png_chunk.data[8]
        self.color_type = self.png_chunk.data[9]
        self.compression_method = self.png_chunk.data[10]
        self.filter_method = self.png_chunk.data[11]
        self.interlace_method = self.png_chunk.data[12]
        print(f"IHDR: 宽： {self.width}, 高： {self.height}")

    def set_source_width(self, width):
        self.width = width

    def set_source_height(self, height):
        self.height = height

    def repair(self):
        """
        修复IHDR chunk
        """
        # 复制chunk的数据，探索症结，修改bit
        chunk_bak = bytearray(self.png_chunk.data)
        # 尝试修改宽高
        n = 4095
        print(f"目标CRC： {self.png_chunk.crc}")
        for w in range(self.width, n):
            for h in range(self.height, n):
                chunk_bak[0:4] = struct.pack(">I", w)
                chunk_bak[4:8] = struct.pack(">I", h)
                new_crc = zlib.crc32(b"IHDR" + chunk_bak) & 0xFFFFFFFF
                if new_crc == self.png_chunk.crc:
                    self.width = w
                    self.height = h
                    self.png_chunk.set_data(chunk_bak)
                    print(f"推荐宽： {self.width}, 高： {self.height}")
                    return True
        return False


class PngRestorer(image_restorer.BaseRestorer):
    def __init__(self):
        pass

    def load_png(self, file_path):
        self.chunks = []
        with open(file_path, "rb") as f:
            self.image_data = f.read()
            # 跳过PNG文件头(8字节)
        offset = 8
        while offset < len(self.image_data):
            chunk = PngChunk(self.image_data, offset)
            self.chunks.append(chunk)
            offset += 12 + chunk.length  # 移动到下一个chunk

    def find_exception_chunks(self):
        """
        寻找异常chunk
        """
        for chunk in self.chunks:
            if chunk.check_crc() == False:
                print(f"Exception chunk: {chunk.get_chunk_type()}")
                # 修复chunk
                self.repair_chunk(chunk)

    def repair_chunk(self, chunk):
        """
        修复chunk
        """
        if chunk.get_chunk_type() == "IHDR":
            # 修复IHDR chunk
            ihdr_chunk_repair = IhdrChunkRestorer(chunk)
            repaired = ihdr_chunk_repair.repair()
            if repaired:
                self.image_data = ihdr_chunk_repair.png_chunk.image_data
                # 重新计算CRC
                self.save_png("test-data/2-repaired.png")

    def save_png(self, file_path):
        """
        保存PNG文件
        """
        with open(file_path, "wb") as f:
            f.write(self.image_data)


if __name__ == "__main__":
    png_repair = PngRestorer()
    png_repair.load_png("test-data/dabai.png")
    png_repair.find_exception_chunks()
